#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Power by Zongsheng Yue 2019-01-07 14:36:55
import numpy
import numpy as np
from cv2 import imread
from skimage import img_as_float
import scipy.io
import time
import sys
import cv2

def im2patch(im, pch_size, stride=1):
    '''
    Transform image to patches.
    Input:
        im: 3 x H x W or 1 X H x W image, numpy format
        pch_size: (int, int) tuple or integer
        stride: (int, int) tuple or integer
    '''
    if isinstance(pch_size, tuple):
        pch_H, pch_W = pch_size
    elif isinstance(pch_size, int):
        pch_H = pch_W = pch_size
    else:
        sys.exit('The input of pch_size must be a integer or a int tuple!')

    if isinstance(stride, tuple):
        stride_H, stride_W = stride
    elif isinstance(stride, int):
        stride_H = stride_W = stride
    else:
        sys.exit('The input of stride must be a integer or a int tuple!')


    C, H, W = im.shape
    num_H = len(range(0, H-pch_H+1, stride_H))
    num_W = len(range(0, W-pch_W+1, stride_W))
    num_pch = num_H * num_W
    pch = np.zeros((C, pch_H*pch_W, num_pch), dtype=im.dtype)
    kk = 0
    for ii in range(pch_H):
        for jj in range(pch_W):
            temp = im[:, ii:H-pch_H+ii+1:stride_H, jj:W-pch_W+jj+1:stride_W]
            pch[:, kk, :] = temp.reshape((C, num_pch))
            kk += 1

    return pch.reshape((C, pch_H, pch_W, num_pch))

def noise_estimate(im, pch_size=8):
    '''
    Implement of noise level estimation of the following paper:
    Chen G , Zhu F , Heng P A . An Efficient Statistical Method for Image Noise Level Estimation[C]// 2015 IEEE International Conference
    on Computer Vision (ICCV). IEEE Computer Society, 2015.
    Input:
        im: the noise image, H x W x 3 or H x W numpy tensor, range [0,1]
        pch_size: patch_size
    Output:
        noise_level: the estimated noise level
    '''

    if im.ndim == 3:
        im = im.transpose((2, 0, 1))
    else:
        im = np.expand_dims(im, axis=0)

    # image to patch
    pch = im2patch(im, pch_size, 3)  # C x pch_size x pch_size x num_pch tensor
    num_pch = pch.shape[3]
    pch = pch.reshape((-1, num_pch))  # d x num_pch matrix


    d = pch.shape[0]
    
    mu = pch.mean(axis=1, keepdims=True)  # d x 1
    X = pch - mu
    sigma_X = np.matmul(X, X.transpose()) / num_pch
    sig_value, _ = np.linalg.eigh(sigma_X)
    sig_value.sort()

    for ii in range(-1, -d-1, -1):
        tau = np.mean(sig_value[:ii])
        if np.sum(sig_value[:ii]>tau) == np.sum(sig_value[:ii] < tau):
            return np.sqrt(tau)

def noise_estimate2(im, level=16, pch_size=49):
    '''
    Implement of noise level estimation of the following paper:
    Chen G , Zhu F , Heng P A . An Efficient Statistical Method for Image Noise Level Estimation[C]// 2015 IEEE International Conference
    on Computer Vision (ICCV). IEEE Computer Society, 2015.
    Input:
        im: the noise image, H x W x 3 or H x W numpy tensor, range [0,1]
        level: 将0-255的灰度平均分成多少个等级
        pch_size: patch_size
    Output:
        noise_level: the estimated noise level
    '''
    sigma_vector = np.array([])
    for i in range(level) :
        condition = (im>i/level) & (im<(i+1)/level)
        im_l = im[ condition ]   
        
        L = len(im_l)
        stride_L = 3
        if L < 10000: # 如果像素数量太少，则不进行估计。预设值为5000
            continue
        # image to patch
        num_L = len(range(0, L-pch_size+1, stride_L))  
        pch = np.zeros((pch_size, num_L), dtype=im.dtype)
        for j in range(num_L):
            temp = im_l[j*stride_L : j*stride_L+pch_size]
            pch[:,j] = temp

        d = pch.shape[0]

        mu = pch.mean(axis=1, keepdims=True)  # d x 1
        X = pch - mu
        sigma_X = np.matmul(X, X.transpose()) / num_L
        sig_value, _ = np.linalg.eigh(sigma_X)
        sig_value.sort()

        sigma_sum = 0
        sigma_num = 0
        for ii in range(-1, -d-1, -1):
            tau = np.mean(sig_value[:ii])     
            if np.isnan(tau):
                continue       
            if np.sum(sig_value[:ii]>tau) == np.sum(sig_value[:ii] < tau):
                sigma_num = sigma_num +1
                sigma_sum = sigma_sum + np.sqrt(tau)
        # 每个灰度范围都可能会有多个sigma，目前采用平均的方法计算sigma；也可以用min、max、中值等方法
        print(f"灰度级范围：{i*256/level}-{(i+1)*256/level}'s sigma: {sigma_sum/sigma_num:.5f}")
                # break
        sigma_vector = np.append(sigma_vector, sigma_sum/sigma_num)
    return sigma_vector


if __name__ == '__main__':
    # im = imread('/usr/noise_estimation/lena.jpg') #, cv2.IMREAD_GRAYSCALE)
    # im = imread('/usr/noise_estimation/I-1.bmp', cv2.IMREAD_GRAYSCALE)
    # im = img_as_float(im)
    # # noise_level = [0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.08, 0.09, 0.10, 0.11, 0.12, 0.13]
    # noise_level = [i*0.01 for i in list(range(1, 20, 1) )]
    # # [数*浮点数 for 数 in 列表]
    # for level in noise_level:
    #     sigma = level 
    #     im_noise = im + np.random.randn(*im.shape) * sigma

    #     start = time.time()
    #     est_level = noise_estimate(im_noise, 8)
    #     end = time.time()
    #     time_elapsed = end -start

    #     str_p = "Time: {0:.4f}, Ture Level: {1:6.4f}, Estimated Level: {2:6.4f}"
    #     print(str_p.format(time_elapsed, level, np.abs(est_level-level)*100/level))

    # exit(0)
    pics = [1,2,3,4,5,6]
    B = list(range(50, 130, 10))
    N = list(range(3, 11, 1))
    F = list(range(10, 90, 10))
    X = np.vstack((B, N, F))
    arguments = ['B', 'N', 'F']

    for i in range(0,1):
        path = '../Experiment/Figure/' + arguments[i] + '/' + arguments[i] + '='
        for j in range(0,8):
            Sigma = []
            start = time.time()
            for k in range(1, 7):
                imgPath = path + str(X[i][j]) + '/Pic_' + str(k) + '.bmp'
                im = imread(imgPath, cv2.IMREAD_GRAYSCALE)
                im = img_as_float(im)

                sigma = noise_estimate(im, 7)

                Sigma.append(sigma)

            end = time.time()
            Time = end - start
            filePath = '../Experiment/Data/' + arguments[i] + '/'+ arguments[i] + '=' + str(X[i][j]) + '/WLS/'
            sigmaPath = filePath + 'Sigma.mat'
            timePath = filePath + 'noiseTime.mat'
            scipy.io.savemat(sigmaPath, {'Sigma': Sigma})
            print(sigmaPath)
            scipy.io.savemat(timePath, {'Time': Time})
            print(timePath)
    # for i in len(name):
    #     str = '../Experiment/Figure/' + name(i) + '/'
    #     for j in pics:
    #         im = imread(str + name(i) + '=', str(i)'B/fig (' + str(i)+').bmp', cv2.IMREAD_GRAYSCALE)
    #         im = im[2900:3400, 3500:4000]  #有条纹的区域
    #         # cv2.imshow('im', im)
    #         # cv2.waitKey(0)
    #         im = img_as_float(im)
    #
    #         start = time.time()
    #         print(f'\n第{i}帧图像：\n')
    #         sigma_vector = noise_estimate2(im, level=16, pch_size=100)
    #         filename = 'sigma_vector' + str(i)+'.mat'
    #         scipy.io.savemat(filename, {'sigma_vector': sigma_vector})
    #         end = time.time()
    #         time_elapsed = end -start

        # str_p = "Time: {0:.4f}, Ture Level: {1:6.4f}, Estimated Level: {2:6.4f}"
        # print(str_p.format(time_elapsed, i, est_level*255))

